
import sys
import os
from datetime import datetime
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                         QGridLayout, QPushButton, QLabel, QScrollArea, QSizePolicy, QFrame)
from PyQt5.QtCore import Qt, QRect, QSize, QPoint
from PyQt5.QtGui import QPixmap, QScreen, QGuiApplication, QColor, QPainter, QPen

class MaskWidget(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.start_pos = None
        self.end_pos = None
        self.is_drawing = False
        
        self.setWindowFlags(Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint | Qt.Tool)
        self.setAttribute(Qt.WA_TranslucentBackground)
        self.setMouseTracking(True)
        self.setCursor(Qt.CrossCursor)
        self.border_color = QColor("#07C160")
        
    def draw_corner_rect(self, painter, pos, size=5):
        painter.fillRect(pos.x() - size, pos.y() - size, size * 2, size * 2, self.border_color)

    def paintEvent(self, event):
        painter = QPainter(self)
        painter.setRenderHint(QPainter.Antialiasing)
        
        mask = QColor(0, 0, 0, 100)
        painter.fillRect(self.rect(), mask)
        
        if self.is_drawing and self.start_pos and self.end_pos:
            select_rect = QRect(self.start_pos, self.end_pos).normalized()
            
            painter.setCompositionMode(QPainter.CompositionMode_Clear)
            painter.fillRect(select_rect, Qt.transparent)
            
            painter.setCompositionMode(QPainter.CompositionMode_SourceOver)
            
            pen = QPen(self.border_color)
            pen.setWidth(2)
            pen.setStyle(Qt.DashLine)
            painter.setPen(pen)
            
            painter.drawRect(select_rect)
            
            self.draw_corner_rect(painter, select_rect.topLeft())
            self.draw_corner_rect(painter, select_rect.topRight())
            self.draw_corner_rect(painter, select_rect.bottomLeft())
            self.draw_corner_rect(painter, select_rect.bottomRight())

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.start_pos = event.pos()
            self.is_drawing = True
            self.end_pos = self.start_pos
            self.update()

    def mouseMoveEvent(self, event):
        if self.is_drawing:
            self.end_pos = event.pos()
            self.update()

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.LeftButton and self.is_drawing:
            self.is_drawing = False
            if self.start_pos and self.end_pos:
                rect = QRect(self.start_pos, self.end_pos).normalized()
                if rect.width() > 0 and rect.height() > 0:
                    self.parent().handle_screenshot(rect)
            self.close()

class PlaceholderWidget(QWidget):
    def __init__(self, size, parent=None):
        super().__init__(parent)
        self.setFixedSize(size, size)
        self.setStyleSheet("""
            QWidget {
                background-color: #f5f5f5;
                border: 1px solid #cccccc;
                border-radius: 4px;
            }
        """)
        
        layout = QVBoxLayout(self)
        layout.setAlignment(Qt.AlignCenter)
        
        label = QLabel("等待添加截图")
        label.setStyleSheet("""
            QLabel {
                color: #999999;
                font-size: 12px;
                border: none;
            }
        """)
        layout.addWidget(label)

class ScreenshotTool(QMainWindow):
    def __init__(self):
        super().__init__()
        self.initUI()
        
        self.screenshot_dir = "screenshots"
        if not os.path.exists(self.screenshot_dir):
            os.makedirs(self.screenshot_dir)
            
        self.thumbnails = []
        self.update_thumbnails()
        
    def initUI(self):
        self.setWindowTitle('截图工具')
        self.setFixedSize(670, 600)
        
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)
        main_layout.setContentsMargins(10, 10, 10, 10)
        
        capture_btn = QPushButton('开始截图', self)
        capture_btn.clicked.connect(self.capture_screen)
        capture_btn.setFixedHeight(30)
        main_layout.addWidget(capture_btn)
        
        # 创建一个带边框的外层容器
        outer_frame = QFrame()
        outer_frame.setFrameStyle(QFrame.Box | QFrame.Raised)
        outer_frame.setLineWidth(2)
        outer_frame.setStyleSheet("""
            QFrame {
                border: 2px solid #666666;
                background-color: white;
                border-radius: 8px;
                padding: 5px;
            }
        """)
        outer_layout = QVBoxLayout(outer_frame)
        outer_layout.setContentsMargins(15, 15, 15, 15)
        
        self.scroll = QScrollArea()
        self.scroll.setWidgetResizable(True)
        self.scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        self.scroll.setVerticalScrollBarPolicy(Qt.ScrollBarAsNeeded)
        self.scroll.setStyleSheet("""
            QScrollArea {
                border: none;
                background-color: transparent;
            }
            QScrollBar:vertical {
                width: 10px;
                background: #f5f5f5;
            }
            QScrollBar::handle:vertical {
                background: #cdcdcd;
                border-radius: 5px;
            }
            QScrollBar::add-line:vertical, QScrollBar::sub-line:vertical {
                height: 0px;
            }
        """)
        
        self.scroll_content = QWidget()
        self.grid_layout = QGridLayout(self.scroll_content)
        self.grid_layout.setSpacing(15)
        self.grid_layout.setContentsMargins(5, 5, 5, 5)
        
        self.scroll_content.setFixedWidth(580)
        
        self.scroll.setWidget(self.scroll_content)
        outer_layout.addWidget(self.scroll)
        
        main_layout.addWidget(outer_frame)
        
    def capture_screen(self):
        self.hide()
        screen = QGuiApplication.primaryScreen()
        self.full_pixmap = screen.grabWindow(0)
        
        self.mask_window = MaskWidget(self)
        desktop = QApplication.desktop()
        screen_geometry = desktop.screenGeometry()
        self.mask_window.setGeometry(screen_geometry)
        self.mask_window.showFullScreen()

    def handle_screenshot(self, rect):
        if rect.isValid():
            cropped = self.full_pixmap.copy(rect)
            timestamp = datetime.now().strftime('%Y%m%d%H%M%S')
            filename = os.path.join(self.screenshot_dir, f'截图_{timestamp}.png')
            cropped.save(filename)
            self.show()
            self.update_thumbnails()
        
    def update_thumbnails(self):
        # 清除现有缩略图
        for thumbnail in self.thumbnails:
            self.grid_layout.removeWidget(thumbnail)
            thumbnail.deleteLater()
        self.thumbnails.clear()
        
        thumbnail_size = 170
        
        # 获取现有截图
        files = []
        if os.path.exists(self.screenshot_dir):
            for f in os.listdir(self.screenshot_dir):
                if f.endswith('.png'):
                    full_path = os.path.join(self.screenshot_dir, f)
                    files.append((full_path, os.path.getmtime(full_path)))
        
        # 按修改时间降序排序
        files.sort(key=lambda x: x[1], reverse=True)
        files = [f[0] for f in files]
        
        # 填充九宫格
        for i in range(9):
            row = i // 3
            col = i % 3
            
            if i < len(files):
                # 显示现有截图
                pixmap = QPixmap(files[i])
                container = QWidget()
                container.setFixedSize(thumbnail_size, thumbnail_size)
                container.setStyleSheet("""
                    background-color: white; 
                    border: 1px solid #cccccc;
                    border-radius: 4px;
                """)
                
                thumbnail = QLabel(container)
                thumbnail.setAlignment(Qt.AlignCenter)
                
                scaled_pixmap = pixmap.scaled(thumbnail_size - 20, thumbnail_size - 20, 
                                            Qt.KeepAspectRatio, 
                                            Qt.SmoothTransformation)
                thumbnail.setPixmap(scaled_pixmap)
                
                thumbnail.setGeometry((thumbnail_size - scaled_pixmap.width()) // 2,
                                    (thumbnail_size - scaled_pixmap.height()) // 2,
                                    scaled_pixmap.width(),
                                    scaled_pixmap.height())
            else:
                # 显示占位符
                container = PlaceholderWidget(thumbnail_size)
            
            self.grid_layout.addWidget(container, row, col)
            self.thumbnails.append(container)

def main():
    app = QApplication(sys.argv)
    tool = ScreenshotTool()
    tool.show()
    sys.exit(app.exec_())

if __name__ == '__main__':
    main()
